import pymysql
import re


class Database:
    def __init__(self, table_name='imagestb'):
        # 验证表名是否合法（只允许字母、数字、下划线）
        if not re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', table_name):
            raise ValueError(f"Invalid table name: {table_name}")

        self.table_name = table_name
        self.connection = pymysql.connect(
            host='localhost',
            user='root',
            password='0520Hlw!@#',
            database='image_db',
            charset='utf8mb4',
            cursorclass=pymysql.cursors.DictCursor
        )

    def insert_image(self, original_image, upload_time):
        try:
            with self.connection.cursor() as cursor:
                sql = f"""
                    INSERT INTO {self.table_name} (original_image, upload_time, status)
                    VALUES (%s, %s, %s)
                """
                print(f"执行SQL: 插入图像数据，时间: {upload_time}")
                cursor.execute(sql, (original_image, upload_time, '处理完成'))
                self.connection.commit()
                print("数据库插入提交成功")
        except Exception as e:
            print(f"数据库插入失败: {str(e)}")
            raise e

    def get_last_insert_id(self):
        """获取最后插入的ID"""
        try:
            with self.connection.cursor() as cursor:
                cursor.execute("SELECT LAST_INSERT_ID() as id")
                result = cursor.fetchone()
                return result['id'] if result else None
        except Exception as e:
            print(f"获取最后插入ID失败: {str(e)}")
            return None

    def get_images(self):
        try:
            with self.connection.cursor() as cursor:
                sql = f"SELECT id, upload_time, status FROM {self.table_name} ORDER BY upload_time DESC LIMIT 10"
                cursor.execute(sql)
                result = cursor.fetchall()
                print(f"查询结果: {result}")
                return result if result else []
        except Exception as e:
            print(f"查询图像列表失败: {str(e)}")
            return []

    def get_images_with_data(self):
        """获取所有图像记录，包括图像数据"""
        try:
            with self.connection.cursor() as cursor:
                sql = f"SELECT id, original_image, upload_time, status FROM {self.table_name} ORDER BY upload_time DESC LIMIT 10"
                cursor.execute(sql)
                result = cursor.fetchall()
                return result if result else []
        except Exception as e:
            print(f"查询图像数据失败: {str(e)}")
            return []

    def get_image_by_id(self, image_id):
        """根据ID获取单个图像记录"""
        try:
            with self.connection.cursor() as cursor:
                sql = f"SELECT id, original_image, upload_time, status FROM {self.table_name} WHERE id = %s"
                cursor.execute(sql, (image_id,))
                result = cursor.fetchone()
                return result
        except Exception as e:
            print(f"查询单个图像失败: {str(e)}")
            return None

    def close(self):
        self.connection.close()